from urdf_parser_py.urdf import URDF

class PositionLimit:
    def __init__(self, lower, upper):
        self.lower = lower
        self.upper = upper

def extract_joint_limits(urdf_file, joint_name_order):
    robot = URDF.from_xml_file(urdf_file)
    limits_dict = {}

    for joint_name in joint_name_order:
        joint = robot.joint_map.get(joint_name)
        if joint is not None:
            limit = joint.limit
            if limit is not None:
                # xjz原 urdf中为无限制
                if limit.lower==0:
                    limit.lower=-5
                if limit.upper==0:
                    limit.upper=5
                limits_dict[joint_name] = PositionLimit(limit.lower, limit.upper)
            else:
                print(f"No limit found for {joint_name}")
        else:
            print(f"Joint {joint_name} not found in URDF")

    return limits_dict

if __name__=="__main__":


    # URDF文件路径
    urdf_file_path = "/home/xjz/catkin_ws/src/unitree_ros/robots/aliengo_description/urdf/aliengo.urdf"

    # 关节名称顺序
    joint_name_order = [
        'FL_hip_joint', 'FL_thigh_joint', 'FL_calf_joint',
        'FR_hip_joint', 'FR_thigh_joint', 'FR_calf_joint',
        'RL_hip_joint', 'RL_thigh_joint', 'RL_calf_joint',
        'RR_hip_joint', 'RR_thigh_joint', 'RR_calf_joint'
    ]

    # 获取关节限制
    position_limits = extract_joint_limits(urdf_file_path, joint_name_order)

    # 打印结果
    for joint_name, limit in position_limits.items():
        print(f"{joint_name}: lower={limit.lower}, upper={limit.upper}")